import cv2
import numpy as np
import matplotlib.pyplot as plt
import os

from bs_img import bs_img_real
from bs_img.bs_cfg_real_5m import *
from bs_img import bs_img_base
from bs_imu import bs_imu
from bs_img.bs_lie_solve import *

import time

import traceback


class BS_Solver:
    def __init__(self):
        sf_img_topic = "/camera_sf/image_dh"
        self.last_plane_pts = np.array([])   #保存上一帧点用于配对
        self.is_first_frame = True                  #判断是否为第一帧


        self.invalid_frame_num_plane = 0      #compute continus invalid frame
        self.invalid_frame_num_drogue = 0
        self.invalid_frame_max = 8     #if invaild frame < invalid_frame_max, output results of the last frame

        #Kalman Filter
        self.Q = np.eye(3)
        self.R = np.eye(3)
        self.x_plane = np.array([np.nan, np.nan, np.nan])
        self.w_plane = 1.5
        self.x_drogue = np.array([np.nan, np.nan, np.nan])
        self.w_drogue = 1.5






    def sf_img_cb(self,img_gray,count):
        _, img_bin = cv2.threshold(img_gray, 70, 0xff, cv2.THRESH_BINARY)
        # 如果相机倒置需要旋转一下
        # lf_img_gray = cv2.rotate(lf_img_gray, cv2.ROTATE_180)
        plane_t_vec = [np.nan]*3
        drogue_t_vec = [np.nan]*3
        drogue_pts = []
        
        gt_pts,aera_list = bs_img_real.img_to_pts(img_bin)
        gt_pts = np.array(gt_pts)
        cnt_plane_pts = [ ]

        len_gt_pts = len(gt_pts)
        if len_gt_pts >= 18:    #噪点太多（可能逆光），为了防止其解算时间太长耽误后面的进度，这里直接跳过
            print("点太多")
            return
        print(len_gt_pts)

        # if self.is_first_frame:   #第一帧图像
        #     drogue_result_dict = bs_img_real.solve_drogue(gt_pts, camK_sf, taper_radius=r_drogue)        
        #     if drogue_result_dict["img_valid"]:
        #         drogue_t_vec = drogue_result_dict["t_vec"]
        #         drogue_pts = drogue_result_dict["drogue_px_calcL"]
        #     else:
        #         drogue_pts = []
        #     drogue_pts = np.array(drogue_pts)
        #     if np.size(drogue_pts) >0:
        #         ind = ~np.isin(gt_pts[:, 0], drogue_pts[:, 0]) | ~np.isin(gt_pts[:, 1], drogue_pts[:, 1])
        #         plane_pts = np.array(gt_pts)[ind]
        #         ind2 = ind == False
        #         plane_aeras = np.array(aera_list)[ind2] 
        #     else:
        #         plane_pts = gt_pts
        #         plane_aeras = np.array(aera_list)
        #     if len(plane_pts)>=5:
        #         plane_result_dict = bs_img_base.solve_plane_pt_new(plane_pts,plane_real_ptL,cam_K=camK_sf,distort_coeffs=distortion_coeffs_sf,plane_area_ptL=plane_aeras)
        #         if plane_result_dict["img_valid"]:
        #             self.is_first_frame = False 
        #             cnt_plane_pts = plane_result_dict["plane_px_calcL"]
        #             plane_t_vec = plane_result_dict["t_vec"].flatten()
        #             plane_quat = bs_imu.rot_vec_to_quat(plane_result_dict["r_vec"].flatten())

        # else:   #除第一帧以外的图像
        #     match_pt_idxs = find_matched_pts(self.last_plane_pts,gt_pts - np.mean(gt_pts,axis=0))
        #     flag_empty = True   #指示匹配是否为空
        #     plane_aeras = []
        #     for match_pt_idx in match_pt_idxs:
        #         if match_pt_idx is not None:
        #             cnt_plane_pts.append(gt_pts[match_pt_idx])
        #             plane_aeras.append(gt_pts[match_pt_idx])
        #             flag_empty = False
        #     if flag_empty:
        #         self.is_first_frame = True   #匹配为空重新开始，视为第一帧
            
        #     # 至少有5个点则解算飞机特征点
        #     if len(cnt_plane_pts) >= 5:    
        #         plane_result_dict = bs_img_base.solve_plane_pt_new(cnt_plane_pts,plane_real_ptL,cam_K=camK_sf,distort_coeffs=distortion_coeffs_sf,plane_area_ptL=plane_aeras)
        #         if plane_result_dict["img_valid"]:
        #             is_first_frame = False 
        #             cnt_plane_pts = plane_result_dict["plane_px_calcL"]
        #             plane_t_vec = plane_result_dict["t_vec"].flatten()
        #             plane_quat = bs_imu.rot_vec_to_quat(plane_result_dict["r_vec"].flatten())
        #         else:
        #             cnt_plane_pts = []
        #     else:
        #         cnt_plane_pts = []
        #     cnt_plane_pts = np.array(cnt_plane_pts)
                
        #     if np.size(cnt_plane_pts) >0:
        #         ind = ~np.isin(gt_pts[:, 0], cnt_plane_pts[:, 0]) | ~np.isin(gt_pts[:, 1], cnt_plane_pts[:, 1])
        #         drogue_pts = np.array(gt_pts)[ind]
        #     else:
        #         drogue_pts = gt_pts
            
        #     drogue_result_dict = bs_img_real.solve_drogue(drogue_pts, camK_sf, taper_radius=r_drogue)        
        #     if drogue_result_dict["img_valid"]:
        #         drogue_t_vec = drogue_result_dict["t_vec"]
        #         drogue_pts = drogue_result_dict["drogue_px_calcL"]

                

        
        # self.last_plane_pts = cnt_plane_pts - np.mean(cnt_plane_pts,axis=0)

        # drogue_pts = np.array(drogue_pts)
        # if np.size(drogue_pts)>0:
        #     x = drogue_pts[:,0]
        #     y = drogue_pts[:,1]
        #     x = x[np.newaxis, :]
        #     y = y[np.newaxis, :]

        #     # x_center,y_center,a,b,theta = bs_img_real.fit_ellipse(x,y)
        #     # # center = (h,k)
        #     # axes_length = (a,b)
        #     # angle = theta/np.pi*180
        #     # # # 计算椭圆曲线上的点
        #     # theta = np.linspace(0, 2*np.pi, 100)  # 角度范围
        #     # x = x_center + a * np.cos(theta-angle)
        #     # y = y_center + b * np.sin(theta-angle)

        #     # # 将椭圆曲线上的点坐标转换为整数
        #     # x = np.round(x).astype(int)
        #     # y = np.round(y).astype(int)

        #     # # 绘制椭圆曲线
        #     # for i in range(len(x)):
        #     #     if x[i]>0 and x[i]<2048 and y[i]>0 and y[i]<2448: 
        #     #         img_bin[y[i], x[i]] = 255

        #     theta = np.linspace(0, 2*np.pi, num=100)  # 定义一个单位圆
        #     if len(drogue_pts) >0:
        #         for i in range(0,len(drogue_pts)):
        #             x = 20*np.cos(theta)+drogue_pts[i,0]
        #             y = 20*np.sin(theta)+drogue_pts[i,1]
        #             plt.fill(x,y,'yellow')
        # theta = np.linspace(0, 2*np.pi, num=100)  # 定义一个单位圆
        # if len(cnt_plane_pts)>0:
        #     for i in range(0,len(cnt_plane_pts)):
        #         x = 20*np.cos(theta)+cnt_plane_pts[i,0]
        #         y = 20*np.sin(theta)+cnt_plane_pts[i,1]
        #         plt.fill(x,y,'red')

        # # 显示图像
        # plt.imshow(img_bin, cmap='gray')
        # plt.axis('off')
        
        # plt.savefig('/home/brightsoul/bs_gd_single/bag/sf_marked/'+str(count)+'.jpg')
        # plt.ion()
        # plt.pause(0.1)
        # plt.close()
        for i in range(0,len(gt_pts)):
            cv2.circle(img_gray,(int(gt_pts[i,0]),int(gt_pts[i,1])),radius=8,color=(255,255,255))
        cv2.namedWindow("1",0)
        cv2.resizeWindow("1",1920,1080)
        cv2.imshow("1",img_gray)
        cv2.waitKey(0)

if __name__ == "__main__":
    bs = BS_Solver()
    folder_path = "/home/brightsoul/Desktop"
    # file_list = os.listdir(folder_path)
    # image_files = [file for file in file_list if file.lower().endswith(('.jpg','.png'))]
    # sorted_files = sorted(image_files)
    # count = 10000
    # # for file in sorted_files:
    # while True:
    #     count = count+1
    #     if count%10 == 0 and count >0:
    #         t1 = time.time()
    #         img_gray = cv2.imread(folder_path+'/'+sorted_files[count],0)
    #         bs.sf_img_cb(img_gray,count)
    #         t2 = time.time()
    #         print(t2-t1)

    #     # if count>2500:
    #     #     break
    try:
        filename = "imagesf3.png"
        img_gray = cv2.imread(folder_path + '/' + filename,0)
        bs.sf_img_cb(img_gray, 0)
    except Exception as e:
        print(traceback.format_exc())